//
//  main.cpp
//  iornn-parser-c
//
//  Created by Phong Le on 27/12/14.
//  Copyright (c) 2014 Phong Le. All rights reserved.
//

#include <stdio.h>

/*
 * Reranker.cpp
 *
 *  Created on: Dec 24, 2014
 *      Author: Phong
 */

#include <string>
#include <iostream>
#include <fstream>
#include <vector>
#include "Matrix.h"
#include "Default.h"
#include "Dictionary.h"
#include "IORNN.h"
#include "IORNNTrainer.h"
#include "Treebank.h"
#include "Depstruct.h"
#include "Deptree.h"
#include "Reranker.h"
#include <regex>
#include <getopt.h>
#include <math.h>
#include <time.h>
#include <stdlib.h>
#include <omp.h>

using namespace std;

//----------------------------- default -----------------------------

int dim = 200;
int completeInside = CMPL_NONE;

// for training
real lambda = 1e-3;
real lambdaL = 1e-6;
real largeRandRange = 0.1;
real smallRandRange = 0.001;

real paramLearningRate = 1e-1;
real learningRateDecay = 0;
real wembLearningRate = 1e-3;
int evalDevStep = 2;
int maxDecreaseCount = 3;

int maxNEpoch = 100;
int batchSize = 50;
int nTop = 10;

int nThreads = 2;

bool countPunc = false;

// for kbest
string kbestScoreSuffix = "mstscores";


// paths
string dataDir = "./data/YM-s/data";
string dicDir = "./data/YM-s/dic";
string wembSource = "glove-100d";
string modelDir = "/tmp";

string trainFile="train.conll";
string devFile="dev.conll";
string kbestDevFile = "dev-20-best-mst2ndorder.conll";
string testFile="test.conll";
string kbestTestFile = "test-20-best-mst2ndorder.conll";

// -------------------------------------------------------------------

void loadParams(string fname) {
    ifstream f(fname);
    string line;
    while (getline(f, line)) {
        vector<string> comps = Utils::splitString(line, "[\t ]\+");
        if (comps.size() != 2) continue;
        
        if (comps[0] == "dim")
            dim = stoi(comps[1]);
        if (comps[0] == "completeInside") {
        	if (comps[1] == "left2right")
        		completeInside = CMPL_LEFT2RIGHT;
        	else if (comps[1] == "right2left")
        		completeInside = CMPL_RIGHT2LEFT;
        	else
        		completeInside = CMPL_NONE;
        }
        
        // for training
        if (comps[0] == "lambda")
            lambda = (real)stod(comps[1]);
        if (comps[0] == "lambdaL")
            lambdaL = (real)stod(comps[1]);
        if (comps[0] == "largeRandRange")
            largeRandRange = (real)stod(comps[1]);
        if (comps[0] == "smallRandRange")
            smallRandRange = (real)stod(comps[1]);
        
        if (comps[0] == "paramLearningRate")
            paramLearningRate = (real)stod(comps[1]);
        if (comps[0] == "learningRateDecay")
            learningRateDecay = (real)stod(comps[1]);
        if (comps[0] == "wembLearningRate")
            wembLearningRate = (real)stod(comps[1]);
        
        if (comps[0] == "maxNEpoch")
            maxNEpoch = stoi(comps[1]);
        if (comps[0] == "batchSize")
            batchSize = stoi(comps[1]);
        if (comps[0] == "nTop")
            nTop = stoi(comps[1]);
        
        if (comps[0] == "maxDecreaseCount")
            maxDecreaseCount = stoi(comps[1]);
        if (comps[0] == "evalDevStep")
            evalDevStep = stoi(comps[1]);
        
        if (comps[0] == "nThreads")
            nThreads = stoi(comps[1]);
        
        // for evaluation
        if (comps[0] == "countPunc")
            countPunc = comps[1] == "true" ?  true : false;
        
        // for kbest
        if (comps[0] == "kbestScoreSuffix")
            kbestScoreSuffix = comps[1];
        
        // paths
        if (comps[0] == "dataDir")
            dataDir = comps[1];
        if (comps[0] == "dicDir")
            dicDir = comps[1];
        if (comps[0] == "wembSource")
            wembSource = comps[1];
        if (comps[0] == "modelDir")
            modelDir = comps[1];
        
        if (comps[0] == "trainFile")
            trainFile = comps[1];
        if (comps[0] == "devFile")
            devFile = comps[1];
        if (comps[0] == "kbestDevFile")
            kbestDevFile = comps[1];
        if (comps[0] == "testFile")
            testFile = comps[1];
        if (comps[0] == "kbestTestFile")
            kbestTestFile = comps[1];
    }
    f.close();
}

vector<Matrix*> loadWordEmbeddings(string fname, Dictionary* vocaDic) {
	cout << "load wordembeddings at " << fname << endl;
	vector<Matrix*> L;
	fstream f(fname, ios::in);
	int nwords, wdim;
	string word;
	f >> nwords;
	f >> wdim;

	L.resize(nwords + vocaDic->size());
	for (int i = 0; i < L.size(); i++)
		L[i] = NULL;

	Matrix* sum = Matrix::zeros(wdim);
	for (int i = 0; i < nwords; i++) {
		f >> word;
		int id = vocaDic->add(word);

		Matrix* we = new Matrix(wdim);
		for (int j = 0; j < wdim; j++)
			f >> we->data[j];
		L[id] = we;
		sum->addi(we);
	}
	f.close();

	L.resize(vocaDic->size());
	if (nwords > 0) sum->muli(1./nwords);
	for (int i = 0; i < L.size(); i++) {
		if (L[i] == NULL) {
			Matrix* we = Matrix::normal(wdim, 0, 1)->addi(sum);
			L[i] = we;
		}
	}

	delete sum;
	return L;
}

void train(IORNN* net) {
    
    // load treebanks
    cout << "load treebanks" << endl;
    Treebank* trainTb = Treebank::load(dataDir+"/"+trainFile, net->vocaDic, net->posDic, net->deprelDic);
    Treebank* devTb = Treebank::load(dataDir+"/"+devFile, net->vocaDic, net->posDic, net->deprelDic);
    KBestTreebank* kDevTb = KBestTreebank::load(dataDir+"/"+kbestDevFile, devTb, net->vocaDic, net->posDic, net->deprelDic);
    
    // shuffle train treebank
    cout << "shuffle train treebank" << endl;
    trainTb->shuffle();
    
    // train
    IORNNTrainer trainer;
    trainer.train(net, trainTb, devTb, kDevTb);
}

void train(string paramFile) {
    loadParams(paramFile);
    
    string path = dicDir+ "/" + wembSource;
    Dictionary* vocaDic = Dictionary::create(path + "/words.lst", TEMPLATE_GLOVE);
    vocaDic->loadBinCode(path + "/word_code.lst");
    
    vector<Matrix*> L = loadWordEmbeddings(path+"/wembs.txt", vocaDic);
    
    Dictionary* posDic = Dictionary::create(dicDir + "/pos.lst", TEMPLATE_NONE);
    Dictionary* deprelDic = Dictionary::create(dicDir+"/deprel.lst", TEMPLATE_NONE);
    
    IORNN* net = new IORNN(dim, vocaDic, posDic, deprelDic, L, completeInside);
    
    train(net);
    
    // load treebanks
    //Treebank* tb = Treebank::load(dataDir+"/"+trainFile, vocaDic, posDic, deprelDic);
    //net->checkGradient(tb);
    //Treebank::free(tb);
    
    Dictionary::free(vocaDic);
    Dictionary::free(posDic);
    Dictionary::free(deprelDic);
    delete net;
}

void parse(string netPath, string testTbPath, string kbestTestTbPath, int kbest, real alpha, string output) {
    cout << "load net " << netPath << endl;
    IORNN* net = IORNN::load(netPath);
    Reranker reranker;
    
    cout << "load treebanks" << endl;
    Treebank* testTb = Treebank::load(testTbPath, net->vocaDic, net->posDic, net->deprelDic);
    KBestTreebank* kTestTb = KBestTreebank::load(kbestTestTbPath, testTb, net->vocaDic, net->posDic, net->deprelDic, kbest);
    
    cout << "parse" << endl;
    reranker.parse(net, testTb, kTestTb, kbest, alpha, output);
}

int main(int argc, char *argv[]) {

	openblas_set_num_threads(1);
	omp_set_num_threads(1);
    srand ( time(NULL) );

    if (argc == 3 && string(argv[1]) == "-train")
        train(string(argv[2]));
    
    else if (argc == 7 && string(argv[1]) == "-parse") {
        string netPath = string(argv[3]);
        string initPath = string(argv[2]);
        int kbest = atoi(argv[4]);
        real alpha = (real)atof(argv[5]);
        string output = string(argv[6]);

        loadParams(initPath);
        string testPath = dataDir + "/" + testFile;
        string kbestPath = dataDir + "/" + kbestTestFile;

        parse(netPath, testPath, kbestPath, kbest, alpha, output);
    }
            
    else {
        cout<< "for training:" << endl << "\t-train init_file" << endl
            << "for parsing:" << endl << "\t-parse init_file net_file kbest alpha output_file" << endl;
    }
}




